package distributed_rate_limiter

import (
	"Limiter/connect/rds"
	"context"
	"errors"
	"github.com/redis/go-redis/v9"
	"log/slog"
	"sync"
	"sync/atomic"
	"time"
)

type WindowTimeLimiter struct {
	period    int
	quota     int
	store     *redis.Conn
	keyPrefix string
	align     bool
	sha       string
	// lock
	rescueLock sync.Mutex
	// redis健康标识
	redisAlive uint32
	// redis监控探测任务标识
	monitorStarted bool
}

// PeriodOption 函数选项模式，它允许在创建对象时通过函数来灵活地设置或配置对象的行为或属性
type PeriodOption func(l *WindowTimeLimiter)

const (
	Error     = 0
	Allowed   = 1
	HitQuota  = 2
	OverQuota = 3

	internalOverQuota = 0
	internalAllowed   = 1
	internalHitQuota  = 2
)

var (
	ErrUnknownCode   = errors.New("unknown status code")
	windowTimeScript = `local limit = tonumber(ARGV[1]) local window = tonumber(ARGV[2]) local current = redis.call("INCRBY", KEYS[1], 1) if current == 1 then redis.call("expire", KEYS[1], window) end if current < limit then return 1 elseif current == limit then return 2 else return 0 end`
)

func NewWindowTimeLimiter(ctx context.Context, period, quota int, client *redis.Client, keyPrefix string,
	opts ...PeriodOption) (*WindowTimeLimiter, error) {
	limiter := &WindowTimeLimiter{
		period:    period,
		quota:     quota,
		store:     client.Conn(),
		keyPrefix: keyPrefix,
	}

	for _, opt := range opts {
		opt(limiter)
	}

	hashes, err := rds.LoadScriptsList(ctx, client, []string{windowTimeScript})
	if err != nil {
		return nil, err
	}
	limiter.sha = hashes[0]

	return limiter, nil
}

func (limiter *WindowTimeLimiter) Take(ctx context.Context, key string) (int, error) {
	resp, err := limiter.store.EvalSha(ctx, limiter.sha, []string{limiter.keyPrefix + key}, limiter.quota, limiter.CalcExpireSeconds()).Result()
	if err != nil {
		limiter.startMonitor()
		return Error, err
	}

	code, ok := resp.(int64)
	if !ok {
		return Error, ErrUnknownCode
	}

	switch code {
	case internalOverQuota:
		return OverQuota, nil
	case internalAllowed:
		return Allowed, nil
	case internalHitQuota:
		return HitQuota, nil
	default:
		return Error, ErrUnknownCode
	}
}

// CalcExpireSeconds 计算当前时间距离下一时间窗口的秒数
func (limiter *WindowTimeLimiter) CalcExpireSeconds() int {
	//判断是否需要对齐，时间窗口为天时，以当地的零点作为时间窗口的起点
	if limiter.align {
		//获取读取时间，并计算时间偏移量
		now := time.Now()
		_, offset := now.Zone()
		//调整时间戳，考虑时区偏移量
		unix := now.Unix() + int64(offset)
		//计算时间窗口距离下一次窗口的秒数
		return limiter.period - int(unix%int64(limiter.period))
	}
	//不需要对齐的直接返回
	return limiter.period
}

func Align() PeriodOption {
	return func(l *WindowTimeLimiter) {
		l.align = true
	}
}

// redis健康探测定时任务
func (limiter *WindowTimeLimiter) waitForRedis() {
	//定时器
	ticker := time.NewTicker(pingInterval)
	// 健康探测成功时回调此函数
	defer func() {
		ticker.Stop()
		limiter.rescueLock.Lock()
		limiter.monitorStarted = false
		limiter.rescueLock.Unlock()
	}()

	for range ticker.C {
		ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(pingInterval))
		hashes, client, err := rds.ReLinkAndLoadScript(ctx, []string{windowTimeScript})
		if err != nil && len(hashes) < 1 {
			cancel()
			continue
		}

		limiter.sha = hashes[0]
		limiter.store = client.Conn()
		atomic.StoreUint32(&limiter.redisAlive, 1)
		slog.Info("Redis returns to normal")
		cancel()
		return
	}
}

// 开启redis健康探测
func (limiter *WindowTimeLimiter) startMonitor() {
	limiter.rescueLock.Lock()
	defer limiter.rescueLock.Unlock()
	// 防止重复开启
	if limiter.monitorStarted {
		return
	}
	// 设置任务和健康标识
	limiter.monitorStarted = true
	atomic.StoreUint32(&limiter.redisAlive, 0)
	// 健康探测
	go limiter.waitForRedis()
}
